from __future__ import annotations

import math
import os
from collections import defaultdict
from enum import Enum

import carla

from driving_gym.simulation.adapter.carla.planner.navigation.global_route_planner import (
    GlobalRoutePlanner,
)


def wpt_tuple(waypoint: carla.Waypoint):
    location = waypoint.transform.location
    return location.x, location.y, location.z


def wpt_dict(waypoint: carla.Waypoint):
    return {
        "road_id": waypoint.road_id,
        "lane_id": waypoint.lane_id,
        "s": round(waypoint.s, 5),
    }


def to_carla_waypoint(carla_map: carla.Map, wpt):
    """Convert various waypoint representations to carla.Waypoint"""
    if isinstance(wpt, carla.Waypoint):
        return wpt

    if isinstance(wpt, (tuple, list)) and len(wpt) == 3:
        # Location tuple (x, y, z)
        loc = carla.Location(x=wpt[0], y=wpt[1], z=wpt[2])
        return carla_map.get_waypoint(loc)

    if isinstance(wpt, dict):
        # XODR dict representation
        waypoints = carla_map.get_waypoint_xodr(
            wpt["road_id"], wpt["lane_id"], wpt["s"]
        )
        return waypoints

    return None


def convert_representation(carla_wp, reference):
    """Convert carla waypoint to same representation as reference"""
    if isinstance(reference, carla.Waypoint):
        return carla_wp

    if isinstance(reference, tuple):
        return wpt_tuple(carla_wp)

    if isinstance(reference, dict):
        return wpt_dict(carla_wp)

    return carla_wp


class RoadType(Enum):
    STRAIGHT = "STRAIGHT"  # Long road with minimal curvature
    CURVED = "CURVED"  # Long road with significant curvature
    SHORT = "SHORT"  # Road segments shorter than threshold
    T_SECTION = "T_SECTION"  # Junction with 6 waypoint pairs
    INTERSECTION = "INTERSECTION"  # Junction with 8 waypoint pairs
    RAMP = "RAMP"  # Junction with 2 waypoint pairs
    OTHER = "OTHER"  # Other types of junctions


class RoadClassifier:
    CURVE_THRESHOLD = 0.1  # radians
    MIN_ROAD_LENGTH = 50  # meters, for distinguishing SHORT roads

    def __init__(self, simulator_map):
        if isinstance(simulator_map, carla.Map):
            self._map = simulator_map
        elif os.path.isfile(simulator_map):
            try:
                with open(simulator_map, "r") as file:
                    self._map = carla.Map(
                        os.path.splitext(os.path.basename(simulator_map))[0],
                        file.read(),
                    )
            except Exception as e:
                raise ValueError(f"Failed to load map from file: {e}")
        elif isinstance(simulator_map, str):
            try:
                self._map = carla.Map("Test", simulator_map)
            except Exception as e:
                raise ValueError(f"Failed to load map from string: {e}")
        else:
            raise ValueError("Invalid map input")

        self._classified_roads = defaultdict(list)
        self._processed_junctions = set()
        self._processed_roads = set()
        self._global_planner = GlobalRoutePlanner(self._map, 2)
        self.graph = self._global_planner._graph

        self._init_classfication()

    def get_classification(
        self,
        string_key: bool = True,
        repr_tuple: bool = False,
        repr_dict: bool = True,
    ):
        """Classify road segments in the CARLA map"""
        if repr_tuple:
            fn = wpt_tuple
        elif repr_dict:
            fn = wpt_dict
        else:
            fn = lambda x: x

        result = {}
        for road_type, segments in self._classified_roads.items():
            result[road_type.value if string_key else road_type] = [
                (fn(start_wp), fn(end_wp)) for start_wp, end_wp in segments
            ]

        return result

    def get_related_segment(self, wpt, road_types=None, string_key: bool = True):
        """Get road segments related to the given waypoint.

        Args:
            wpt: Waypoint representation (tuple of (x,y,z), dict with xodr info, or carla.Waypoint)
            road_types: List of RoadType or None (if None, return all types)
            string_key: If True, use RoadType string value as key in output dictionary

        Returns:
            Dictionary of {RoadType: [(start_wp, end_wp), ...]} using same representation as input
        """
        # Convert input waypoint to carla.Waypoint for internal processing
        carla_wpt = to_carla_waypoint(self._map, wpt)
        if carla_wpt is None:
            return {}

        # If road_types not specified, use all types
        if road_types is None:
            road_types = list(RoadType)
        elif isinstance(road_types, (RoadType, str)):
            road_types = [road_types]
        road_types = [RoadType(rt) for rt in road_types]

        related_segments = defaultdict(list)

        # Helper function to check if waypoint is related to segment
        def is_related(wp, start_wp, end_wp):
            # Same road ID
            if wp.road_id == start_wp.road_id or wp.road_id == end_wp.road_id:
                return True

            # Connected through junction
            if start_wp.is_junction or end_wp.is_junction:
                junction_id = (
                    start_wp.get_junction().id
                    if start_wp.is_junction
                    else end_wp.get_junction().id
                )
                if wp.is_junction and wp.get_junction().id == junction_id:
                    return True
                # Check if waypoint is connected to junction
                next_wps = wp.next(2.0)
                prev_wps = wp.previous(2.0)
                for next_wp in next_wps + prev_wps:
                    if next_wp.is_junction and next_wp.get_junction().id == junction_id:
                        return True
            return False

        # Find related segments
        for road_type in road_types:
            key = road_type.value if string_key else road_type
            for start_wp, end_wp in self._classified_roads[road_type]:
                if is_related(carla_wpt, start_wp, end_wp):
                    # Convert waypoints to same representation as input
                    converted_segment = (
                        convert_representation(start_wp, wpt),
                        convert_representation(end_wp, wpt),
                    )
                    related_segments[key].append(converted_segment)

        return dict(related_segments)

    def _init_classfication(self):
        self._process_all_roads()
        self._process_junctions()

    def _process_all_roads(self):
        """Process all roads in the map using road_id and lane_id"""
        # Get all waypoints in the map with a reasonable sampling distance
        all_waypoints = self._map.generate_waypoints(2.0)

        # Group waypoints by road_id and lane_id
        road_lanes = defaultdict(list)
        for wp in all_waypoints:
            road_lanes[(wp.road_id, wp.lane_id)].append(wp)

        # Process each road
        for (road_id, lane_id), waypoints in road_lanes.items():
            if (road_id, lane_id) in self._processed_roads:
                continue

            # Get complete waypoint sequence for this road segment
            sample_wp = waypoints[0]
            prev_wpts = sample_wp.previous_until_lane_start(2)
            next_wpts = sample_wp.next_until_lane_end(2)

            for wp in prev_wpts[::-1]:
                if len(prev_wpts) == 1 or (
                    wp.road_id == sample_wp.road_id and wp.lane_id == sample_wp.lane_id
                ):
                    start_wp = wp
                    break
            for wp in next_wpts[::-1]:
                if len(next_wpts) == 1 or (
                    wp.road_id == sample_wp.road_id and wp.lane_id == sample_wp.lane_id
                ):
                    end_wp = wp
                    break

            # Get all waypoints between start and end
            full_path = self._get_complete_waypoint_path(start_wp, end_wp)

            # Calculate road length
            road_length = self._calculate_path_length(full_path)

            if road_length < self.MIN_ROAD_LENGTH:
                self._classified_roads[RoadType.SHORT].append((start_wp, end_wp))
            else:
                # Classify based on curvature
                road_type = self._classify_road_curvature(full_path)
                self._classified_roads[road_type].append((start_wp, end_wp))

            self._processed_roads.add((road_id, lane_id))

    def _process_junctions(self):
        for edge in self.graph.edges:
            edge_data = self.graph.edges[edge]
            start_wp = edge_data["entry_waypoint"]
            end_wp = edge_data["exit_waypoint"]

            if edge_data["intersection"]:
                # Skip if we've already processed this junction
                junction = start_wp.get_junction()
                if junction.id in self._processed_junctions:
                    continue

                # Get all junction waypoints
                junction_waypoints = junction.get_waypoints(carla.LaneType.Driving)
                junction_type = self._classify_junction_type(junction_waypoints)

                # Add to classified roads
                self._classified_roads[junction_type].append((start_wp, end_wp))
                self._processed_junctions.add(junction.id)

    def _get_complete_waypoint_path(self, start_wp, end_wp):
        """Get complete waypoint path between start and end waypoints"""
        path = [start_wp]
        current_wp = start_wp

        while True:
            next_wps = current_wp.next(2.0)
            if not next_wps or current_wp.road_id != next_wps[0].road_id:
                break
            else:
                current_wp = next_wps[0]
                path.append(current_wp)

        path.append(end_wp)
        return path

    def _calculate_path_length(self, waypoints):
        """Calculate total path length through waypoints"""
        length = 0.0
        for i in range(len(waypoints) - 1):
            length += waypoints[i].transform.location.distance(
                waypoints[i + 1].transform.location
            )
        return length

    def _classify_junction_type(self, junction_waypoints):
        """Classify junction type based on number of roads connected"""
        distinct_roads = set()
        for start, end in junction_waypoints:
            distinct_roads.add(start.road_id)
            distinct_roads.add(end.road_id)

        num_roads = len(distinct_roads)
        if 2 <= num_roads <= 4:
            return RoadType.RAMP
        elif num_roads == 6:
            return RoadType.T_SECTION
        elif num_roads >= 12:
            return RoadType.INTERSECTION
        else:
            return RoadType.OTHER

    def _classify_road_curvature(self, waypoints):
        """Classify road segment as STRAIGHT or CURVED based on waypoint path"""
        max_angle_deviation = 0

        for i in range(len(waypoints) - 1):
            curr_wp = waypoints[i]
            next_wp = waypoints[i + 1]

            # Get forward vectors
            curr_vec = curr_wp.transform.rotation.get_forward_vector()
            next_vec = next_wp.transform.rotation.get_forward_vector()

            # Calculate angle between vectors
            dot_product = (
                curr_vec.x * next_vec.x
                + curr_vec.y * next_vec.y
                + curr_vec.z * next_vec.z
            )
            dot_product = max(min(dot_product, 1.0), -1.0)
            angle = abs(math.acos(dot_product))

            max_angle_deviation = max(max_angle_deviation, angle)

        return (
            RoadType.CURVED
            if max_angle_deviation > self.CURVE_THRESHOLD
            else RoadType.STRAIGHT
        )

    def get_statistics(self):
        """Returns count of each road type"""
        return {
            road_type: len(segments)
            for road_type, segments in self._classified_roads.items()
        }
